package midend.IR_Optimizer;

import midend.IR_Instruction.Calc_Instruction;
import midend.IR_Instruction.Instruction;
import midend.IR_Value.IR_BaseBlock;
import midend.IR_Value.IR_Int_Value;
import midend.IR_Value.IR_Float_Value;
import midend.IR_Value.IR_Local_Value;
import midend.IR_Value.IR_Value;
import midend.IR_Type.IR_Int32;
import midend.IR_Type.IR_Float;

import java.util.*;

import static midend.IR_Type.IR_Const.*;

public class ArithmeticOptimizer {

    private IR_BaseBlock baseBlock;
    private LinkedList<Instruction> instructions;
    private Map<String, ArithmeticPattern> patterns;
    private boolean optimized = false;

    public ArithmeticOptimizer(IR_BaseBlock baseBlock) {
        this.baseBlock = baseBlock;
        this.instructions = new LinkedList<>();
        this.patterns = new HashMap<>();
    }

    /**
     * 表示算术运算模式的内部类
     */
    private static class ArithmeticPattern {
        String baseValue;      // 基础值（如 %a21）
        IR_Value base;
        String operation;      // 操作类型（ADD, SUB, MUL, DIV, SREM等）
        int count;            // 重复次数
        int othercount;
        boolean isFloat;      // 是否为浮点运算
        List<String> resultNames; // 中间结果名称列表
        List<IR_Value> resultValue; // 中间结果名称列表

        ArithmeticPattern(String baseValue, String operation, boolean isFloat,IR_Value base) {
            this.baseValue = baseValue;
            this.operation = operation;
            this.isFloat = isFloat;
            this.count = 0;
            this.othercount = 0;
            this.resultNames = new ArrayList<>();
            this.resultValue = new ArrayList<>();
            this.base = base;
        }

        void addOccurrence(String resultName,IR_Value value,String optype) {
            if(this.operation.equals(optype)){
                this.count++;
            }
            else{
                othercount++;
            }
            this.resultNames.add(resultName);
            this.resultValue.add(value);
        }
    }

    /**
     * 主要的优化入口
     */
    public boolean optimize() {
        // 这里假设从baseBlock中获取指令列表
        // 实际实现中需要根据你的框架结构调整
        collectInstructions();

        // 识别算术模式
        identifyArithmeticPatterns();

        // 执行优化
        return applyOptimizations();
    }

    /**
     * 收集基本块中的指令
     */
    private void collectInstructions() {
        // 根据你的框架结构，这里需要调整为从baseBlock获取指令的实际方法
        // 假设有一个方法可以获取指令列表
        // this.instructions = baseBlock.getInstructions();
        instructions.clear();
        instructions.addAll(baseBlock.get_Instruction_List());
    }

    /**
     * 识别算术运算模式
     */
    private void identifyArithmeticPatterns() {
        patterns.clear();

        for (Instruction instr : instructions) {
            if (instr instanceof Calc_Instruction) {
                Calc_Instruction calcInstr = (Calc_Instruction) instr;
                analyzeCalcInstruction(calcInstr);
            }
        }
    }

    /**
     * 分析单个计算指令
     */
    private void analyzeCalcInstruction(Calc_Instruction calcInstr) {
        String opType = calcInstr.get_op_Type();
        List<IR_Value> operands = calcInstr.get_op();

        if (operands.size() != 3) return; // 结果 = 操作数1 操作 操作数2

        IR_Value result = operands.get(0);
        IR_Value op1 = operands.get(1);
        IR_Value op2 = operands.get(2);

        // 检查是否为重复的算术运算模式
        if (isRepeatedArithmeticPattern(opType, op1, op2)) {
            handleRepeatedPattern(opType, op1, op2, result);
        }
    }

    /**
     * 检查是否为重复的算术运算模式
     */
    private boolean isRepeatedArithmeticPattern(String opType, IR_Value op1, IR_Value op2) {
        // 支持的可优化运算
        if (!Arrays.asList(ADD, SUB,FADD, FSUB).contains(opType)) {
            return false;
        }

        // 检查是否为 x + x, x - x, x * x 等模式
        if (op1.get_ValueName().equals(op2.get_ValueName())) {
            return true;
        }

        // 检查是否为 result + baseValue 的链式模式
        if (isChainedPattern(op1, op2)) {
            return true;
        }

        return false;
    }

    /**
     * 检查是否为链式模式（如 %add25 + %a21）
     */
    private boolean isChainedPattern(IR_Value op1, IR_Value op2) {
        String op1Name = op1.get_ValueName();
        String op2Name = op2.get_ValueName();

        // 检查是否存在已有模式，其中op1是之前的结果，op2是基础值
        for (ArithmeticPattern pattern : patterns.values()) {
            if (pattern.resultNames.contains(op1Name) && pattern.baseValue.equals(op2Name)) {
                return true;
            }
        }

        return false;
    }

    /**
     * 处理重复模式
     */
    private void handleRepeatedPattern(String opType, IR_Value op1, IR_Value op2, IR_Value result) {
        String op1Name = op1.get_ValueName();
        String op2Name = op2.get_ValueName();
        String resultName = result.get_ValueName();

        boolean isFloat = result.get_Type() instanceof IR_Float;

        // 情况1: x + x 模式
        if (op1Name.equals(op2Name)) {
            String patternKey = op1Name + "_" + opType;
            ArithmeticPattern pattern = patterns.get(patternKey);
            if (pattern == null) {
                pattern = new ArithmeticPattern(op1Name, opType, isFloat,op1);
                patterns.put(patternKey, pattern);
            }
            pattern.addOccurrence(resultName,result,opType);
        }

        // 情况2: 链式模式
        else if (isChainedPattern(op1, op2)) {
            // 找到对应的模式并更新
            for (ArithmeticPattern pattern : patterns.values()) {
                if (pattern.resultNames.contains(op1Name) && pattern.baseValue.equals(op2Name)) {
                    pattern.addOccurrence(resultName,result,opType);
                    break;
                }
            }
        }
    }

    /**
     * 应用优化
     */
    private boolean applyOptimizations() {
        boolean hasOptimized = false;

        for (ArithmeticPattern pattern : patterns.values()) {
            if (pattern.count > 1) {
                hasOptimized |= optimizePattern(pattern);
            }
        }

        return hasOptimized;
    }

    /**
     * 优化单个模式
     */
    private boolean optimizePattern(ArithmeticPattern pattern) {
        switch (pattern.operation) {
            case ADD:
            case FADD:
                return optimizeAddition(pattern);
            case SUB:
            case FSUB:
                return optimizeSubtraction(pattern);
            default:
                return false;
        }
    }

    /**
     * 优化加法运算
     * 例如: %a21 + %a21 + %a21 + ... (n次) = %a21 * n
     */
    private boolean optimizeAddition(ArithmeticPattern pattern) {
        if (pattern.count < 2) return false;

        // 创建乘法指令替换多个加法
        String finalResult = pattern.resultNames.get(pattern.resultNames.size() - 1);
        IR_Value finalvalue = pattern.resultValue.get(pattern.resultValue.size() - 1);

        // 创建常数值
        IR_Value multiplier;
        if (pattern.isFloat) {
            multiplier = new IR_Float_Value(pattern.count + 1 - pattern.othercount, IR_Float.IR_Float(), true);
        } else {
            multiplier = new IR_Int_Value(pattern.count + 1 - pattern.othercount, IR_Int32.IR_Int32(), true);
        }

        // 创建优化后的乘法指令
        Calc_Instruction optimizedInstr = createOptimizedInstruction(
                finalvalue,
                pattern.base,
                multiplier,
                pattern.isFloat ? FMUL : MUL
        );

        // 替换原有指令
        replaceInstructions(pattern, optimizedInstr);

        return true;
    }

    /**
     * 优化减法运算
     * 例如: %a21 - %a21 - %a21 - ... = %a21 * (1 - n)
     */
    private boolean optimizeSubtraction(ArithmeticPattern pattern) {
        if (pattern.count < 2) return false;

        String finalResult = pattern.resultNames.get(pattern.resultNames.size() - 1);
        IR_Value finalvalue = pattern.resultValue.get(pattern.resultValue.size() - 1);

        // 计算系数 (1 - n)
        int coefficient = 1 - pattern.count + pattern.othercount;

        IR_Value multiplier;
        if (pattern.isFloat) {
            multiplier = new IR_Float_Value(coefficient, IR_Float.IR_Float(), true);
        } else {
            multiplier = new IR_Int_Value(coefficient, IR_Int32.IR_Int32(), true);
        }

        Calc_Instruction optimizedInstr = createOptimizedInstruction(
                finalvalue,
                pattern.base,
                multiplier,
                pattern.isFloat ? FMUL : MUL
        );

        replaceInstructions(pattern, optimizedInstr);

        return true;
    }


    /**
     * 创建优化后的指令
     */
    private Calc_Instruction createOptimizedInstruction(IR_Value result, IR_Value baseValue, IR_Value multiplier, String opType) {

        List<IR_Value> operands = Arrays.asList(result, baseValue, multiplier);

        return new Calc_Instruction(operands, baseBlock, opType);
    }

    /**
     * 替换指令
     */
    private void replaceInstructions(ArithmeticPattern pattern, Calc_Instruction optimizedInstr) {
        // 1. 移除原有的重复指令
        // 2. 插入优化后的指令
        // 3. 更新数据流和控制流


        // 找到需要替换的指令
        List<Instruction> toRemove = new ArrayList<>();
        for (String resultName : pattern.resultNames) {
            for (Instruction instr : instructions) {
                if (instr instanceof Calc_Instruction) {
                    Calc_Instruction calcInstr = (Calc_Instruction) instr;
                    if (calcInstr.get_op().get(0).get_ValueName().equals(resultName)) {
                        toRemove.add(instr);
                        break;
                    }
                }
            }
        }

        Instruction lastinstr = null;
        lastinstr = toRemove.get(toRemove.size() - 1);
        int index = instructions.indexOf(lastinstr);
        // 添加新指令
        instructions.add(index + 1,optimizedInstr);
        // 移除旧指令
        instructions.removeAll(toRemove);

        optimized = true;
    }

    public LinkedList<Instruction> instructions() {
        return instructions;
    }
}